# coding = utf-8

'''
训练3d的dense u-net
'''

from network.HDenseUnet import denseUnet3dForTrain,denseUnet
from dataset.IRCAD_HDENSE_UNET import IRCAD_DATA
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
import torch
from tqdm import tqdm
import numpy as np

def train(data_loader, criterion, optimizer, unet_2d_model, unet_3d_model, num_slide):
    loss_list = []
    tbar = tqdm(data_loader, ascii=True, desc='train', dynamic_ncols=True)

    for batch_idx,(data,mask) in enumerate(tbar):
        data = data.cuda()
        mask = mask.cuda()
        input2d = data[:, 0:2, :, :]
        single = data[:, 0:1, :, :]
        input2d = torch.cat((input2d, single), 1)
        for i in range(num_slide - 2):
            input2dtmp = data[:, i:i + 3, :, :]
            input2d = torch.cat((input2d, input2dtmp), 0)
            if i == num_slide - 3:
                f1 = data[:, num_slide - 2: num_slide, :, :]
                f2 = data[:, num_slide - 1: num_slide, :, :]
                ff = torch.cat((f1, f2), 1)
                input2d = torch.cat((input2d, ff), 0)

        with torch.no_grad():
            (feature, cls) = unet_2d_model(input2d)

        cls = cls.clone().permute(1, 0, 2, 3)
        cls.unsqueeze_(0)
        feature = feature.clone().permute(1,0,2,3)
        feature.unsqueeze_(0)

        x_tmp = data.clone().unsqueeze(0)
        x_tmp *= 250.0

        input3d = torch.cat((cls, x_tmp), 1)

        optimizer.zero_grad()
        output = unet_3d_model(input3d, feature)
        loss = criterion(output, mask)
        loss.backward()
        optimizer.step()

        loss_list.append(loss.item())

        tbar.set_postfix(loss=f'{loss.item():.5f}')

    return loss_list

def train3d_unet():
    ircad = IRCAD_DATA(origion_data_path="/datasets/3Dircadb/origion", hdense_data_path="/datasets/3Dircadb/HDenseNet",
                       input_size=224, input_cols=8, mean=0)
    subset = ircad.train_dataset
    data_loader = DataLoader(subset, batch_size=1, shuffle=True)
    unet_2d_model = torch.load("/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/dense_unet_2d/epoll_357.pkl").cuda()
    unet_3d_model = denseUnet3dForTrain().cuda()
    #unet_3d_model = torch.load("/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/dense_unet_3d/epoll_186.pkl").cuda()
    opt = torch.optim.SGD(unet_3d_model.parameters(), lr=1e-3, momentum=0.9)
    loss = torch.nn.CrossEntropyLoss(weight=torch.tensor((0.78, 0.65, 8.57), device='cuda'), reduction='mean')
    model_save_path = "/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/dense_unet_3d"

    for i in range(400):
        print("-------------------epoll{}---------------------------".format(i))

        unet_3d_model.train()
        unet_2d_model.eval()
        loss_list = train(data_loader=data_loader, criterion=loss, optimizer=opt, unet_2d_model=unet_2d_model, unet_3d_model=unet_3d_model,
                          num_slide=8)
        torch.save(unet_3d_model, "{}/epoll_{}.pkl".format(model_save_path, str(i).zfill(3)))
        print()
        print("train loss : {}".format(np.mean(np.array(loss_list))))
        print()


if __name__ == '__main__':
    train3d_unet()